"""
定义：优化问题的通用状态定义方法
"""
import numpy as np
from opt_func.FuncSet import * #目前先测试该函数
import copy

class Same_range_problems(object):
    def __init__(self, **kwargs):
        self.range_lower = float(kwargs.get('lower', -2))
        self.range_upper = float(kwargs.get('upper', 2))
        self.history_steps = 3 #回溯前3步路径
        self.dim_num = int(kwargs.get('dim_num', 3))
        # 变量池：暂时用顺序的数字序号表示，（之后代码可能是按特定顺序优化的）
        self.dim_rank = np.linspace(0,self.dim_num-1,self.dim_num)
        self.split_num = int(kwargs.get('split_num', 5)) #每一维度水平数，暂时都相同
        self.func = eval(str(kwargs.get('func_name', 'Rosenbrock')))
        # 每个维度变量区间，可以修改


    def init_problem(self):
        # 变量池：暂时用顺序的数字序号表示，（之后代码可能是按特定顺序优化的）
        # self.avaliable_move_idx = list(np.linspace(0,self.value_num-1,self.value_num))

        # self.avaliable_move_value = np.round(np.repeat(
        #     np.linspace(self.range_lower, self.range_upper, self.split_num)[np.newaxis, :], self.dim_num, axis=0), 2)
        self.avaliable_move_value = np.round(np.linspace(self.range_lower,
                                                         self.range_upper,
                                                         self.split_num), 2)

        self.last_move_idx = -1
        self.states = {} # key:value_idx, value:value {x1:value,x2:value}

    def current_state(self):
        square_state = np.zeros((self.history_steps, self.dim_num, self.split_num))
        if self.states:
            value_idxes, values = np.array(list(zip(*self.states.items())))
            locations=[location for value_idx in value_idxes for (location, option) in enumerate(self.dim_rank)
              if value_idx==option]
            level_location = [location for value in values for (location, option) in enumerate(self.avaliable_move_value)
              if value==option]

            len_state = len(self.states)
            square_state[0][locations, level_location] = 1.

            if len_state > 2:
                square_state[1][locations[:-1], level_location[:-1]] = 1.
                square_state[2][locations[:-2], level_location[:-2]] = 1.
            elif len_state==2:
                square_state[1][locations[:-1], level_location[:-1]] = 1.


        return square_state

    def do_move(self, move):
        """
        :param move: 数值move_value
        :return:
        """
        # move_idx, move_value = list(zip(*move))
        # self.avaliable_move_idx.remove(move_idx)
        self.last_move_idx += 1  # 按顺序（即0,1，2...）调整到下一个变量
        self.states[self.last_move_idx] = move # 更新states

    def get_avaliable_move(self):
        next_move_idx = self.last_move_idx + 1
        return next_move_idx, self.avaliable_move_value

    def end_optimization(self):
        if len(self.states) == self.dim_num:
            target = self.func(np.array(list(self.states.values()))) #问题的真实目标值
            ## 需要进行目标值向当前路径价值的转化
            # 将target数值控制在-100~100
            # target_upper = 1000
            # target_lower = 0
            # if target > target_upper:
            #     target = target_upper
            # elif target < target_lower:
            #     target = target_lower
            # # 将target数值映射为[-1, 1]区间
            # value = -1 + (1-(-1))/(target_upper-target_lower)*(target-target_lower)
            # value = -value #由于目标为最小值，因此取负数
            value = -np.log(target) #对预测值取对苏，消除残差的正态性
            return True, value
        else:
            return False, None

class Game():

    def __init__(self, same_range_problem):
        self.srp = same_range_problem

    def start_self_play(self, player, temp=1e-3):
        self.srp.init_problem()
        states, mcts_probs, func_targets = [], [], []
        while True:
            value, value_prob = player.get_action(self.srp, temp=temp, return_prob=1)
            mcts_probs.append(value_prob)
            self.srp.do_move(value)
            states.append(self.srp.current_state())

            end, fun_target = self.srp.end_optimization()
            if end:
                # print(('游戏结束,函数目标值:{}, 参数为:{}').format(
                #                                         fun_target,
                #                                         states[-1][0]))
                func_targets = [fun_target] * len(states)
                return list(zip(states, mcts_probs, func_targets))

    def start_play(self, player, temp=1e-3):
        self.srp.init_problem()
        while True:
            value = player.get_action(self.srp, temp=temp, return_prob=0)
            self.srp.do_move(value)
            end, fun_target = self.srp.end_optimization()
            if end:
                state = self.srp.current_state()[0]
                # print(('测试结束,函数目标值:{}, 参数为:{}').format(
                #                                         fun_target,
                #                                         state))
                return fun_target, state



if __name__ == '__main__' :
    srp = Same_range_problems()
    srp.init_problem()
